# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule
import torch.nn.functional as F

from mmseg.ops import resize
from ..builder import HEADS
from .aspp_head import ASPPHead
from .sep_aspp_head import DepthwiseSeparableASPPModule

class ProjectionHead(nn.Module):
    """
    The projection head used by contrast learning.
    Args:
        dim_in (int): The dimensions of input features.
        proj_dim (int, optional): The output dimensions of projection head. Default: 256.
        proj (str, optional): The type of projection head, only support 'linear' and 'convmlp'. Default: 'convmlp'.
    """
    def __init__(self, dim_in, proj_dim=256, proj='convmlp'):
        super(ProjectionHead, self).__init__()
        
        
        if proj == 'linear':
            self.proj = nn.Conv2d(dim_in, proj_dim, kernel_size=1)
        elif proj == 'convmlp':
            
            self.proj = nn.Sequential(
                nn.Conv2d(dim_in, dim_in, kernel_size=1),
                nn.BatchNorm2d(dim_in),
                nn.ReLU(),
                nn.Conv2d(dim_in, proj_dim, kernel_size=1)
            )

    def forward(self, x):
        return F.normalize(self.proj(x), p=2, dim=1)


@HEADS.register_module()
class DepthwiseSeparableASPPHeadContrastive(ASPPHead):
    """Encoder-Decoder with Atrous Separable Convolution for Semantic Image
    Segmentation.

    This head is the implementation of `DeepLabV3+
    <https://arxiv.org/abs/1802.02611>`_.

    Args:
        c1_in_channels (int): The input channels of c1 decoder. If is 0,
            the no decoder will be used.
        c1_channels (int): The intermediate channels of c1 decoder.
    """

    def __init__(self, c1_in_channels, c1_channels, **kwargs):
        super(DepthwiseSeparableASPPHeadContrastive, self).__init__(**kwargs)
        assert c1_in_channels >= 0
        self.aspp_modules = DepthwiseSeparableASPPModule(
            dilations=self.dilations,
            in_channels=self.in_channels,
            channels=self.channels,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        if c1_in_channels > 0:
            self.c1_bottleneck = ConvModule(   # 为低级特征进行降维
                c1_in_channels,
                c1_channels,
                1,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg)
        else:
            self.c1_bottleneck = None
        self.sep_bottleneck = nn.Sequential(  # 最后处理concat的特征
            DepthwiseSeparableConvModule(
                self.channels + c1_channels,
                self.channels,
                3,
                padding=1,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg),
            DepthwiseSeparableConvModule(
                self.channels,
                self.channels,
                3,
                padding=1,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg))
        
        self.proj_head = ProjectionHead(dim_in=self.channels,proj_dim=256)
        
        # self.unsample = nn.Sequential(
        #     nn.ConvTranspose2d(
        #     self.channels,
        #     self.channels,
        #     kernel_size=4, 
        #     stride=2, 
        #     padding=1),
        #     nn.ConvTranspose2d(
        #     self.channels,
        #     self.channels,
        #     kernel_size=4, 
        #     stride=2, 
        #     padding=1)
        # )
    
    def forward(self, inputs):
        """Forward function."""
        # print(inputs[0].shape)  # 6,256,128,128
        # print(inputs[1].shape)  # 6,512,64,64
        # print(inputs[2].shape)  # 6,1024,64,64
        # print(inputs[3].shape) # 6,2048,64,64
        x = self._transform_inputs(inputs)
        # print(x.shape)   # 6,2048,64,64  取得最后一个
        # print(embedding.shape)   # 6,256,64,64
        

        # add done 
        aspp_outs = [
            resize(
                self.image_pool(x),  # 先池化再降维  6,512,1,1
                size=x.size()[2:],
                mode='bilinear',
                align_corners=self.align_corners)
        ]
        # print(aspp_outs[0].shape)  # 6,512,64,64
        aspp_outs.extend(self.aspp_modules(x))
        aspp_outs = torch.cat(aspp_outs, dim=1)
        output = self.bottleneck(aspp_outs)
        # print("output:" , output.shape)  # 6,512,64,64
        if self.c1_bottleneck is not None:
            c1_output = self.c1_bottleneck(inputs[0])  # 低级特征过1x1卷积为通道降维
            # print("c1_output" , c1_output.shape)  # 6,48,128,128  256-->48
            output = resize(
                input=output,
                size=c1_output.shape[2:],
                mode='bilinear',
                align_corners=self.align_corners)
            # print("output:" , output.shape) # 6,512,128,128
            output = torch.cat([output, c1_output], dim=1)  # 语义和细节concat
            # print("after concat output:" , output.shape)  # 6,560,128,128
        output = self.sep_bottleneck(output)  # 最后过两个3*3深度可分离卷积
        # print("after sep_bottleneck output:" , output.shape)  # 6,512,128,128
        
        ### 先上采样
        # output = self.unsample(output)  # 6 , 512, 512,512
        output_size = output.shape[2:]
        # output_size = output_size * 2 
        output_size = torch.tensor(output_size) * 2 
        # print(output_size)
        output = resize(
            output,  
            size=torch.Size(output_size),
            mode='bilinear',
            align_corners=self.align_corners
        )
        output_size = output.shape[2:]
        output_size = torch.tensor(output_size) * 2 
        output = resize(
            output,  
            size=torch.Size(output_size),
            mode='bilinear',
            align_corners=self.align_corners
        )
        embedding = None
        if hasattr(self,'proj_head'):
            embedding = self.proj_head(output)   
        output = self.cls_seg(output)
        # print("final output" , output.shape)   # 6,2,512,512
        return {'out' : output , 'embed' : embedding}
